/*
 * MIT License
 *
 * Copyright (c) 2020 wen.gu <454727014@qq.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

/***************************************************************************
* Name: av_decoder.cpp
*
* Purpose: decoder component implementation with libav
*
* Developer:
*   wen.gu , 2018-12-28
*
* TODO:
*
***************************************************************************/

/******************************************************************************
**    INCLUDES
******************************************************************************/
#define _CRT_SECURE_NO_WARNINGS
#include <Windows.h>
#include <stdio.h>
#include <string>

#include "mp_plugin.h"
#include "mp_plugin_manager.h"

#include "av_decoder.h"


#define LOG_TAG "avDecoder"
#include "mp_log.h"

/******************************************************************************
**    MACROS
******************************************************************************/

#define AVD_COMP_NAME "av decoder"
#define AVD_COMP_VER "00.01.00"

/******************************************************************************
**    VARIABLE DEFINITIONS
******************************************************************************/


/******************************************************************************
**   FUNCTION DEFINITIONS
******************************************************************************/
AvDecoder::AvDecoder()
    :MPDecoder(AVD_COMP_NAME, AVD_COMP_VER)
{
    /** todo somethings */
    avcodec_register_all();
}

AvDecoder::~AvDecoder()
{
    /** todo somethings */
}

///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////

MPError AvDecoder::reset() /** reset decoder state to new decode */
{
    MPError ret = MP_ErrorInvalidStatus;
    if (mCtx)
    {
        avcodec_flush_buffers(mCtx);
        ret = MP_OK;
    }
    else
    {
        LOGE("decodet ctx is NULL, cann't reset\n");
    }

    return ret;
}

MPTrackInfo* AvDecoder::getFormat()/**get decoder output information*/
{
    return &mInfo;
}

MPError AvDecoder::sendPacket(MPBuffer* pkt)
{
    MPError ret = MP_ErrorBadParameter;

    if (pkt)
    {   
        AVPacket avp = { 0 };
        int  got_frame = 0;
        av_init_packet(&avp);

        avp.data = pkt->data();
        avp.size = pkt->size();
        avp.dts = pkt->dts();
        avp.pts = pkt->pts();

        /* send the packet with the compressed data to the decoder */
        int res = avcodec_send_packet(mCtx, &avp);
        if (res == 0)
        {
            ret = MP_OK;
        }
        else
        {
            LOGE("Error submitting the packet to the decoder\n");
            ret = MP_ErrorUndefined;
        }

    }

    return ret;
}

MPError AvDecoder::receiveFrame(MPBuffer** frame)
{
    MPError ret = MP_ErrorBadParameter;
    if (frame)
    {
        AVFrame avframe = { 0 };

        int res = avcodec_receive_frame(mCtx, &avframe);
        if (res == 0)
        {
            if (mInfo.category == eMediaCategory::Video)
            {
                ret = processVideoOutput(avframe, frame);
            }
            else if (mInfo.category == eMediaCategory::Audio)
            {
                ret = processAudioOutput(avframe, frame);
            }
            else
            {
                LOGE("not support this category: %d\n", mInfo.category);
                ret = MP_ErrorUndefined;
            }
        }
        else if ((res == AVERROR(EAGAIN)) || (res == AVERROR_EOF))
        {
            ret = MP_ErrorEof;
        }        
        else if (res < 0) 
        {
            LOGE("Error during decoding\n");
            ret = MP_ErrorUndefined;
        }
    }

    return ret;
}

///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////

MPError AvDecoder::initialize(const MPTrackInfo& info)
{
    MPError ret = MP_ErrorUnsupportedType;
    AVCodecID avid = convert2avId(info.format);

    if (avid != AV_CODEC_ID_NONE)
    {
        mCodec = avcodec_find_decoder(avid);

        if (mCodec)
        {
            mCtx = avcodec_alloc_context3(mCodec);

            if (mCtx)
            {
                ret = MP_OK;
                if (mCodec->capabilities & CODEC_CAP_TRUNCATED)
                {
                    mCtx->flags |= CODEC_CAP_TRUNCATED;
                }

                if (mCodec->type == AVMEDIA_TYPE_VIDEO)
                {
                    const MPVideoInfo& vt = (info.video);
                    mCtx->width = vt.picWidth;
                    mCtx->height = vt.picHeight;
                    mCtx->framerate.num = vt.framerate;
                    mCtx->framerate.den = 1;

                }
                else if (mCodec->type == AVMEDIA_TYPE_AUDIO)
                {
                    const MPAudioInfo& at = (info.audio);

                    mCtx->sample_rate = at.samplerate;
                    mCtx->channels = at.channelCount;
                    mCtx->bits_per_coded_sample = at.bitDepth;
                }
                else
                {
                    LOGE("unsuport this category(%d, avmedia:%d)\n", info.category, mCodec->type);
                    ret = MP_ErrorUnsupportedType;
                }

                if (MP_OK == ret)
                {    
                    mInfo = info;

                    if (info.extraData && (info.extraSize > 0))
                    {
                        mInfo.extraData = MAKE_SHARED_ARRAY8(info.extraSize);
                        memcpy(mInfo.extraData.get(), info.extraData.get(), info.extraSize);
                    }
                    mCtx->bit_rate = mInfo.bit_rate;
                    mCtx->extradata = mInfo.extraData.get();
                    mCtx->extradata_size = mInfo.extraSize;
                    int res = avcodec_open2(mCtx, mCodec, nullptr);

                    if (res == 0)
                    {

                    }
                    else
                    {
                        LOGE("open avcodec failed(%d)\n", res);
                        ret = MP_ErrorUndefined;
                    }
                }
            }
            else
            {
                LOGD("Could not allocate video codec context\n");
            }
        }
        else
        {
            LOGD("Could not open codec\n");
        }
    }

    return ret;


}

///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////

AVCodecID AvDecoder::convert2avId(eMediaFormat fmt)
{
    AVCodecID ret = AV_CODEC_ID_NONE;
    switch (fmt)
    {
    case eMediaFormat::AAC:    ret = AV_CODEC_ID_AAC; break;
    case eMediaFormat::Vorbis: ret = AV_CODEC_ID_VORBIS; break;
    case eMediaFormat::AVC:    ret = AV_CODEC_ID_H264; break;
    case eMediaFormat::MP3:    ret = AV_CODEC_ID_MP3; break;
    default:
        break;
    }

    return ret;
}

eMediaFormat AvDecoder::convert2Fmt(AVPixelFormat id)
{
    eMediaFormat ret = eMediaFormat::Unknown;
    switch (id)
    {
    case AV_PIX_FMT_NV12: ret = eMediaFormat::YUV420SemiPlanar; break;
    case AV_PIX_FMT_YUV420P: ret = eMediaFormat::YUV420Planar; break;
    default: break;
    }

    return ret;
}

eMediaFormat AvDecoder::convert2Fmt(AVSampleFormat id)
{
    eMediaFormat ret = eMediaFormat::Unknown;
    switch (id)
    {
    case AV_SAMPLE_FMT_S16:  ret = eMediaFormat::S16L; break;
    case AV_SAMPLE_FMT_S32:  ret = eMediaFormat::S32L; break;
    default: break;
    }

    return ret;
}


void AvDecoder::audioPlanar2Interleaved(uint8_t* dst, AVFrame& frame)
{
    uint32_t i = 0;
    uint32_t j = 0;
    uint8_t *psrc = frame.data[0];
    uint32_t sample_cnt = frame.nb_samples;
    uint32_t sample_size = av_get_bytes_per_sample((AVSampleFormat)frame.format);
    uint32_t channel_cnt = frame.channels;
    uint32_t seg_size = sample_size / channel_cnt; /** bytes per sample per channel */

    for (i = 0; i < channel_cnt; i++)
    {
        uint8_t* pdst = dst;
        uint8_t* psrc = frame.data[i];

        for (j = 0; j < sample_cnt; j++)
        {
            memcpy(pdst, psrc, seg_size);
            pdst += sample_size;
            psrc += seg_size;
        }

        dst += seg_size;
    }
}

uint32_t AvDecoder::getFrameSize(AVFrame& avf)
{
    uint32_t frameSize = avf.width * avf.height;
    switch (avf.format)
    {
    case AV_PIX_FMT_YUV420P:
    case AV_PIX_FMT_NV12:
        frameSize = frameSize * 3 / 2;
        break;   
    default: frameSize = 0; break;
    }

    return frameSize;
}

MPError AvDecoder::processAudioOutput(AVFrame& avf, MPBuffer** frame)
{
    MPError ret = MP_OK;
    uint32_t frame_size = av_samples_get_buffer_size(NULL, avf.channels, avf.nb_samples, (AVSampleFormat)avf.format, 1);
    uint32_t width = avf.width;
    uint32_t height = avf.height;
    *frame = new MPBuffer(frame_size);

    BufferInfo& bi = (*frame)->info();

    if (av_sample_fmt_is_planar((AVSampleFormat)avf.format)) /** planar audio */
    {
        audioPlanar2Interleaved((*frame)->data(), avf);
    }
    else /** interleaved audio */
    {
        memcpy((*frame)->data(), avf.data[0], avf.linesize[0]);
    }

    bi.format = convert2Fmt((AVSampleFormat)avf.format);
    bi.pts = av_frame_get_best_effort_timestamp(&avf);
    bi.dts = -1;
    bi.category = eMediaCategory::Audio;
    bi.audio.bitDepth = av_get_bytes_per_sample((AVSampleFormat)avf.format) * 8;
    bi.audio.channelCount = avf.channels;
    bi.audio.samplerate = avf.sample_rate;

    return ret;
}

MPError AvDecoder::processVideoOutput(AVFrame& avf, MPBuffer** frame)
{
    MPError ret = MP_OK;
    uint32_t line_size = avf.linesize[0];
    uint8_t* psrc = avf.data[0];
    uint32_t frameSize = getFrameSize(avf);

    if (frameSize > 0)
    {
        *frame = new MPBuffer(frameSize);
        uint8_t* pdst = (*frame)->data();
        uint32_t width = avf.width;
        uint32_t height = avf.height;
        BufferInfo& bi = (*frame)->info();

        bi.format = convert2Fmt((AVPixelFormat)avf.format);
        bi.pts = av_frame_get_best_effort_timestamp(&avf);
        bi.dts = -1;
        bi.category = eMediaCategory::Video;
        bi.video.picWidth = width;
        bi.video.picHeight = height;   

        if (width == line_size) /** process Y(luma) */
        {
            uint32_t luma_size = width * height;
            memcpy(pdst, psrc, luma_size);
            pdst += luma_size;

            switch (avf.format)
            {
            case AV_PIX_FMT_NV12:         
                memcpy(pdst, avf.data[1], luma_size / 2);
                break;
            case  AV_PIX_FMT_YUV420P:
                memcpy(pdst, avf.data[1], luma_size / 4);
                pdst += (luma_size / 4);
                memcpy(pdst, avf.data[2], luma_size / 4);
                break;
            default:
                ret = MP_ErrorUndefined;
                break;
            }
        }
        else
        {
            uint32_t i = 0;
            uint32_t stride = avf.linesize[0];
            for (i = 0; i < height; i++)
            {
                memcpy(pdst, psrc, width);
                pdst += width;
                psrc += stride;
            }

            switch (avf.format)
            {
            case AV_PIX_FMT_NV12:
                psrc = avf.data[1];
                for (i = 0; i < height; i++)
                {
                    memcpy(pdst, psrc, width / 2);
                    pdst += width / 2;
                    psrc += stride;
                }

                break;
            case  AV_PIX_FMT_YUV420P:
                psrc = avf.data[1];
                for (i = 0; i < height / 2; i++)
                {
                    memcpy(pdst, psrc, width / 2);
                    pdst += width / 2;
                    psrc += stride;
                }

                psrc = avf.data[2];
                for (i = 0; i < height / 2; i++)
                {
                    memcpy(pdst, psrc, width / 2);
                    pdst += width / 2;
                    psrc += stride;
                }
                break;
            default:
                ret = MP_ErrorUndefined;
                break;
            }
        }
    }
    else
    {
        LOGE("invalid framesize: 0 \n");
        ret = MP_ErrorUnsupportedType;
    }

    return ret;
}

///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////
class AvDecoderPlugin : public MPPlugin
{
public:
    AvDecoderPlugin(const std::string& name,
        const std::string& version,
        const std::string& author
    );
    virtual ~AvDecoderPlugin();

public:
    virtual MPError createComponent(const MPCCreateParam* param, MPComponent** comp) override;
};

///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////
AvDecoderPlugin::AvDecoderPlugin(const std::string& name,
    const std::string& version,
    const std::string& author)
    :MPPlugin(name, version, author)
{
    /** todo something */
}

AvDecoderPlugin::~AvDecoderPlugin()
{
    /** todo something */
}

///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////

MPError AvDecoderPlugin::createComponent(const MPCCreateParam* param, MPComponent** comp)
{
    MPError ret = MP_ErrorBadParameter;

    if (param && comp)
    {
        if (param->type == MP_DECODER_COMP_TYPE)
        {
            const MPDecoderCreateParam* dp = static_cast<const MPDecoderCreateParam*>(param);

            if (dp)
            {
                AvDecoder* avd = new AvDecoder;

                ret = avd->initialize(dp->info);

                if (MP_OK == ret)
                {
                    *comp = avd;
                }
                else
                {
                    delete avd;
                }
            }
        }
        else
        {
            ret = MP_ErrorUnsupportedType;
        }
    }

    return ret;
}

////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////

/** the entry of current plugin */
MP_PLUGIN_ENTRY(AvDecoderPlugin,
    "av decoder plugin",  /** name */
    "00.01.00",           /** version */
    "454727014@qq.com" /** author */
);
